/**
 * $RCSfile$
 * $Revision: $
 * $Date: $
 *
 * Copyright (C) 2005-2008 Jive Software. All rights reserved.
 *
 * This software is published under the terms of the GNU Public License (GPL),
 * a copy of which is included in this distribution, or a commercial license
 * agreement with Jive.
 */

package org.xmppbot.net;

import org.dom4j.DocumentHelper;
import org.dom4j.Element;
import org.dom4j.Namespace;
import org.dom4j.QName;
import org.jivesoftware.openfire.lockout.LockOutManager;
import org.jivesoftware.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xmppbot.Connection;
import org.xmppbot.session.IncomingServerSession;
import org.xmppbot.session.LocalSession;
import org.xmppbot.session.Session;

import javax.security.sasl.SaslException;
import javax.security.sasl.SaslServer;
import java.io.UnsupportedEncodingException;
import java.util.*;

/**
 * SASLAuthentication is responsible for returning the available SASL mechanisms to use and for
 * actually performing the SASL authentication.<p>
 * <p/>
 * The list of available SASL mechanisms is determined by:
 * <ol>
 * <li>Whether the underlying connection has been secured or not.</li>
 * </ol>
 *
 * TODO done refactoring
 *
 * @author Hao Chen
 * @author Gaston Dombiak
 */
class SASLAuthentication {

    private static final Logger logger = LoggerFactory.getLogger(SASLAuthentication.class);

    /**
     * The utf-8 charset for decoding and encoding Jabber packet streams.
     */
    protected static String CHARSET = "UTF-8";

    private static final String SASL_NAMESPACE = "xmlns=\"urn:ietf:params:xml:ns:xmpp-sasl\"";

    private static Map<String, ElementType> typeMap = new TreeMap<String, ElementType>();

    private static Set<String> mechanisms = new HashSet<String>() {{
            mechanisms.add("ANONYMOUS");
            mechanisms.add("EXTERNAL");
    }};

    public enum ElementType {

        AUTH("auth"), RESPONSE("response"), CHALLENGE("challenge"), FAILURE("failure"), UNDEF("");

        private String name = null;

        public String toString() {
            return name;
        }

        private ElementType(String name) {
            this.name = name;
            typeMap.put(this.name, this);
        }

        public static ElementType valueof(String name) {
            if (name == null) {
                return UNDEF;
            }
            ElementType e = typeMap.get(name);
            return e != null ? e : UNDEF;
        }
    }

    @SuppressWarnings({"UnnecessarySemicolon"})
    // Support for QDox Parser
    public enum Status {
        /**
         * Entity needs to respond last challenge. Session is still negotiating
         * SASL authentication.
         */
        needResponse,
        /**
         * SASL negotiation has failed. The entity may retry a few times before the connection
         * is closed.
         */
        failed,
        /**
         * SASL negotiation has been successful.
         */
        authenticated;
    }

    /**
     * Returns a string with the valid SASL mechanisms available for the specified session. If
     * the session's connection is not secured then only include the SASL mechanisms that don't
     * require TLS.
     *
     * @param session The current session
     * @return a string with the valid SASL mechanisms available for the specified session.
     */
    public static String getSASLMechanisms(LocalSession session) {
        if (!(session instanceof IncomingServerSession)) {
            return "";
        }
        StringBuilder sb = new StringBuilder(195);
        sb.append("<mechanisms xmlns=\"urn:ietf:params:xml:ns:xmpp-sasl\">");
        if (session instanceof IncomingServerSession) {
            // Server connections dont follow the same rules as clients
            if (session.isSecure()) {
                boolean usingSelfSigned = false;
//                Certificate[] certificates = session.getConnection().getLocalCertificates();
//                for (Certificate certificate : certificates) {
//                    try {
//                        if (CertificateManager
//                                .isSelfSignedCertificate(SSLConfig.getKeyStore(), (X509Certificate) certificate)) {
//                            usingSelfSigned = true;
//                        }
//                    } catch (Exception e) {
//                        usingSelfSigned = true;
//                    }
//                }
                if (!usingSelfSigned) {
                    // Offer SASL EXTERNAL only if TLS has already been negotiated and we are not
                    // using a self-signed certificate
                    sb.append("<mechanism>EXTERNAL</mechanism>");
                }
            }
        } else {
            for (String mech : getSupportedMechanisms()) {
                sb.append("<mechanism>");
                sb.append(mech);
                sb.append("</mechanism>");
            }
        }
        sb.append("</mechanisms>");
        return sb.toString();
    }

    public static Element getSASLMechanismsElement(Session session) {
        if (!(session instanceof IncomingServerSession)) {
            return null;
        }

        Element mechs = DocumentHelper.createElement(new QName("mechanisms",
                new Namespace("", "urn:ietf:params:xml:ns:xmpp-sasl")));
        if (session instanceof IncomingServerSession) {
            // Server connections dont follow the same rules as clients
            if (session.isSecure()) {
                // Offer SASL EXTERNAL only if TLS has already been negotiated
                Element mechanism = mechs.addElement("mechanism");
                mechanism.setText("EXTERNAL");
            }
        } else {
            for (String mech : getSupportedMechanisms()) {
                Element mechanism = mechs.addElement("mechanism");
                mechanism.setText(mech);
            }
        }
        return mechs;
    }

    /**
     * Handles the SASL authentication packet. The entity may be sending an initial
     * authentication request or a response to a challenge made by the server. The returned
     * value indicates whether the authentication has finished either successfully or not or
     * if the entity is expected to send a response to a challenge.
     *
     * @param session the session that is authenticating with the server.
     * @param doc     the stanza sent by the authenticating entity.
     * @return value that indicates whether the authentication has finished either successfully
     *         or not or if the entity is expected to send a response to a challenge.
     * @throws java.io.UnsupportedEncodingException
     *          If UTF-8 charset is not supported.
     */
    public static Status handle(LocalSession session, Element doc) throws UnsupportedEncodingException {
        logger.debug("SASLAuthentication: {} {}", session, doc);
        Status status;
        String mechanism;
        if (doc.getNamespace().asXML().equals(SASL_NAMESPACE)) {
            ElementType type = ElementType.valueof(doc.getName());
            switch (type) {
                case AUTH:
                    mechanism = doc.attributeValue("mechanism");
                    // Store the requested SASL mechanism by the client
                    session.setSessionData("SaslMechanism", mechanism);
                    //logger.debug("SASLAuthentication.doHandshake() AUTH entered: "+mechanism);
                    if (mechanism.equalsIgnoreCase("ANONYMOUS") &&
                            mechanisms.contains("ANONYMOUS")) {
                        status = doAnonymousAuthentication(session);
                    } else if (mechanism.equalsIgnoreCase("EXTERNAL")) {
                        status = doExternalAuthentication(session, doc);
                    } else {
                        logger.warn("Client wants to do a MECH we don't support: '" +
                                mechanism + "'");
                        authenticationFailed(session);
                        status = Status.failed;
                    }
                    break;
                case RESPONSE:
                    // Store the requested SASL mechanism by the client
                    mechanism = (String) session.getSessionData("SaslMechanism");
                    if (mechanism.equalsIgnoreCase("EXTERNAL")) {
                        status = doExternalAuthentication(session, doc);
                    } else if (mechanisms.contains(mechanism)) {
                        SaslServer ss = (SaslServer) session.getSessionData("SaslServer");
                        if (ss != null) {
                            boolean ssComplete = ss.isComplete();
                            String response = doc.getTextTrim();
                            try {
                                if (ssComplete) {
                                    authenticationSuccessful(session, ss.getAuthorizationID(),
                                            null);
                                    status = Status.authenticated;
                                } else {
                                    byte[] data = StringUtils.decodeBase64(response);
                                    if (data == null) {
                                        data = new byte[0];
                                    }
                                    byte[] challenge = ss.evaluateResponse(data);
                                    if (ss.isComplete()) {
                                        authenticationSuccessful(session, ss.getAuthorizationID(),
                                                challenge);
                                        status = Status.authenticated;
                                    } else {
                                        // Send the challenge
                                        sendChallenge(session, challenge);
                                        status = Status.needResponse;
                                    }
                                }
                            }
                            catch (SaslException e) {
                                logger.debug("SASLAuthentication: SaslException", e);
                                authenticationFailed(session);
                                status = Status.failed;
                            }
                        } else {
                            logger.error("SaslServer is null, should be valid object instead.");
                            authenticationFailed(session);
                            status = Status.failed;
                        }
                    } else {
                        logger.warn(
                                "Client responded to a MECH we don't support: '" + mechanism + "'");
                        authenticationFailed(session);
                        status = Status.failed;
                    }
                    break;
                default:
                    authenticationFailed(session);
                    status = Status.failed;
                    // Ignore
                    break;
            }
        } else {
            logger.debug("SASLAuthentication: Unknown namespace sent in auth element: " + doc.asXML());
            authenticationFailed(session);
            status = Status.failed;
        }
        // Check if SASL authentication has finished so we can clean up temp information
        if (status == Status.failed || status == Status.authenticated) {
            // Remove the SaslServer from the Session
            session.removeSessionData("SaslServer");
            // Remove the requested SASL mechanism by the client
            session.removeSessionData("SaslMechanism");
        }
        return status;
    }


    private static Status doAnonymousAuthentication(LocalSession session) {
        // TODO check anonymous allowed
        if (false) {
            // Verify that client can connect from his IP address
            boolean forbidAccess = false;
//            try {
//                String hostAddress = session.getConnection().getHostAddress();
//                if (!LocalClientSession.getAllowedAnonymIPs().isEmpty() &&
//                        !LocalClientSession.getAllowedAnonymIPs().containsKey(hostAddress)) {
//                    byte[] address = session.getConnection().getAddress();
//                    String range1 = (address[0] & 0xff) + "." + (address[1] & 0xff) + "." +
//                            (address[2] & 0xff) +
//                            ".*";
//                    String range2 = (address[0] & 0xff) + "." + (address[1] & 0xff) + ".*.*";
//                    String range3 = (address[0] & 0xff) + ".*.*.*";
//                    if (!LocalClientSession.getAllowedAnonymIPs().containsKey(range1) &&
//                            !LocalClientSession.getAllowedAnonymIPs().containsKey(range2) &&
//                            !LocalClientSession.getAllowedAnonymIPs().containsKey(range3)) {
//                        forbidAccess = true;
//                    }
//                }
//            } catch (UnknownHostException e) {
//                forbidAccess = true;
//            }
            if (forbidAccess) {
                authenticationFailed(session);
                return Status.failed;
            }
            // Just accept the authentication :)
            authenticationSuccessful(session, null, null);
            return Status.authenticated;
        }
        else {
            // anonymous login is disabled so close the connection
            authenticationFailed(session);
            return Status.failed;
        }
    }

    private static Status doExternalAuthentication(LocalSession session, Element doc)
            throws UnsupportedEncodingException {
        // At this point the connection has already been secured using TLS

        if (session instanceof IncomingServerSession) {

            String hostname = doc.getTextTrim();
            if (hostname == null || hostname.length() == 0) {
                // No hostname was provided so send a challenge to get it
                sendChallenge(session, new byte[0]);
                return Status.needResponse;
            }

            hostname = new String(StringUtils.decodeBase64(hostname), CHARSET);
            // Check if cerificate validation is disabled for s2s
            // Flag that indicates if certificates of the remote server should be validated.
            // Disabling certificate validation is not recommended for production environments.
            boolean verify = true;
//                    JiveGlobals.getBooleanProperty("xmpp.server.certificate.verify", true);
            if (!verify) {
                authenticationSuccessful(session, hostname, null);
                return Status.authenticated;
            }
            // Check that hostname matches the one provided in a certificate
            Connection connection = session.getConnection();

//            for (Certificate certificate : connection.getPeerCertificates()) {
//                for (String identity : CertificateManager.getPeerIdentities((X509Certificate) certificate)) {
//                    // Verify that either the identity is the same as the hostname, or for wildcarded
//                    // identities that the hostname ends with .domainspecified or -is- domainspecified.
//                    if ((identity.startsWith("*.")
//                            && (hostname.endsWith(identity.replace("*.", "."))
//                            || hostname.equals(identity.replace("*.", ""))))
//                            || hostname.equals(identity)) {
//                        authenticationSuccessful(session, hostname, null);
//                        return Status.authenticated;
//                    }
//                }
//            }

        } else {
            logger.debug("SASLAuthentication: unknown session type. Cannot perform EXTERNAL authentication");
        }
        authenticationFailed(session);
        return Status.failed;
    }

    private static void sendChallenge(Session session, byte[] challenge) {
        StringBuilder reply = new StringBuilder(250);
        if (challenge == null) {
            challenge = new byte[0];
        }
        String challenge_b64 = StringUtils.encodeBase64(challenge).trim();
        if ("".equals(challenge_b64)) {
            challenge_b64 = "="; // Must be padded if null
        }
        reply.append(
                "<challenge xmlns=\"urn:ietf:params:xml:ns:xmpp-sasl\">");
        reply.append(challenge_b64);
        reply.append("</challenge>");
        session.deliver(reply.toString());
    }

    private static void authenticationSuccessful(LocalSession session, String username,
                                                 byte[] successData) {
        if (username != null && LockOutManager.getInstance().isAccountDisabled(username)) {
            // Interception!  This person is locked out, fail instead!
            LockOutManager.getInstance().recordFailedLogin(username);
            authenticationFailed(session);
            return;
        }
        StringBuilder reply = new StringBuilder(80);
        reply.append("<success xmlns=\"urn:ietf:params:xml:ns:xmpp-sasl\"");
        if (successData != null) {
            String successData_b64 = StringUtils.encodeBase64(successData).trim();
            reply.append(">").append(successData_b64).append("</success>");
        } else {
            reply.append("/>");
        }
        session.deliver(reply.toString());
        String hostname = username;
        // Add the validated domain as a valid domain. The remote server can
        // now send packets from this address
//        ((LocalIncomingServerSession) session).addValidatedDomain(hostname);
    }

    private static void authenticationFailed(LocalSession session) {
        StringBuilder reply = new StringBuilder(80);
        reply.append("<failure xmlns=\"urn:ietf:params:xml:ns:xmpp-sasl\">");
        reply.append("<not-authorized/></failure>");
        session.deliver(reply.toString());
        // Give a number of retries before closing the connection
        Integer retries = (Integer) session.getSessionData("authRetries");
        if (retries == null) {
            retries = 1;
        } else {
            retries = retries + 1;
        }
        session.setSessionData("authRetries", retries);
//        if (retries >= JiveGlobals.getIntProperty("xmpp.auth.retries", 3)) {
        if (retries >= 3) {
            // Close the connection
            session.close();
        }
    }

    /**
     * Adds a new SASL mechanism to the list of supported SASL mechanisms by the server. The
     * new mechanism will be offered to clients and connection managers as stream features.<p>
     * <p/>
     * Note: this method simply registers the SASL mechanism to be advertised as a supported
     * mechanism by Openfire. Actual SASL handling is done by Java itself, so you must add
     * the provider to Java.
     *
     * @param mechanism the new SASL mechanism.
     */
    public static void addSupportedMechanism(String mechanism) {
        mechanisms.add(mechanism);
    }

    /**
     * Removes a SASL mechanism from the list of supported SASL mechanisms by the server.
     *
     * @param mechanism the SASL mechanism to remove.
     */
    public static void removeSupportedMechanism(String mechanism) {
        mechanisms.remove(mechanism);
    }

    /**
     * Returns the list of supported SASL mechanisms by the server. Note that Java may have
     * support for more mechanisms but some of them may not be returned since a special setup
     * is required that might be missing. Use {@link #addSupportedMechanism(String)} to add
     * new SASL mechanisms.
     *
     * @return the list of supported SASL mechanisms by the server.
     */
    public static Set<String> getSupportedMechanisms() {
        return new HashSet<String>(mechanisms);
    }

}